from numpy import *
from scipy import *
from matplotlib.pyplot import *
from scipy.misc import derivative
from FunctionalBohm import Schrodinger1D
from AtomicUnits import hbar # we set hbar == 1.0 for convenience...

#controlling array size, advisable to have large for decent results
N=1000
#declare arrays zero to avoid problems
x=zeros((N+4))
V=zeros((N+4))
k1=zeros((N+4))
k2=zeros((N+4))
k3=zeros((N+4))
k4=zeros((N+4))
Psi_superpos=zeros((N+4))

#declare variables and constants (consider stability when choosing dx, dt)
t=0.0; t_end=4000; dt=0.001; sizingconst=1.0; dx= 0.0001#sizingconst/N
mass=2.0; omega=3.0

#initial conditions
i=array(range(-2,N+2))
x=(sizingconst*i)/N
V=0.5*mass**2*omega**2*x**2

def f(xarg, targ,i): #define function to be solved

    def Psi1(xarg): # the n=1 eigenstate
        normalisation = 1.0/sqrt(2.0) * (mass*omega/(hbar*pi))**0.25
        gauss = exp(-mass*omega*xarg[i]**2/(2.0*hbar))
        hermitearg = xarg[i]*sqrt(mass*omega/hbar)
        def hermite1(y):
            return 2.0 * y
        time = exp((0-1.0j)*omega*(1.5)*targ)
        return normalisation*gauss*hermite1(hermitearg)*time

    def Psi2(xarg): # the n=2 eigenstate
        normalisation = 1.0/sqrt(8.0) * (mass*omega/(hbar*pi))**0.25
        gauss = exp(-mass*omega*xarg[i]**2/(2.0*hbar))
        hermitearg = xarg[i]*sqrt(mass*omega/hbar)
        def hermite2(y):
            return 4.0 * y**2 - 2.0
        time = exp((0-1.0j)*omega*(2.5)*targ)
        return normalisation*gauss*hermite2(hermitearg)*time

    def Psi_superpos(xarg): # the superposition (equal contributions)
        return 1/sqrt(2) * (Psi1(xarg) + Psi2(xarg))


    Sch = Schrodinger1D(V,Psi_superpos,mass,dx=1.0)# NOTE trying changing this to avoid bug
    return (1/mass)*Sch.get_momentum()(xarg)



#RK4 method

def RK4(f,t,x,i,dt):
    
    k1[i]=dt*f(x,t,i)            
    k2[i]=dt*f(x+0.5*k1 , t+0.5*dt, i)        
    k3[i]=dt*f(x+0.5*k2, t+0.5*dt, i)          
    k4[i]=dt*f(x+k3, t+dt, i)           

    x=(1.0/6.0)*(k1+2.0*k2+2.0*k3+k4)
    p=mass*f(x,t,i)
    return x,p#(1.0/6.0)*(k1+2.0*k2+2.0*k3+k4)   # RK4 equation

X=[]
T=[]
P=[]
#X.append(x)
#T.append(t)

for t in range (0,t_end):

    x_,p_=RK4(f,t,x,i,dt)
    x=x+x_
    t=t+dt

    #saving results to array
    X.append(x)
    T.append(t)
    P.append(p_)

X_=array(X)
T_=array(T)
P_=array(P)

#singling out trajectories to plot
m=arange(200, 800, 25)
X_plot=X_*1000
Y=X_plot[:, m]
Z=P_[:, m]

#plotting options

#plot(T_,X_)
plot(T_,Y)
#plot(T_,P_)
#plot(T_,Z)
#scatter(X_,P_)
#scatter(Y,Z)
    
print("Updated")
